#
# Copyright (c) 2011-2012 Johann Hanne
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#


class OMF_Record(object):
    def __init__(self, rawcontent):
        self.rectype = ord(rawcontent[0])
        self.reclen = ord(rawcontent[1]) | (ord(rawcontent[2]) << 8)

        self.checksum = 0
        for c in rawcontent[0:self.reclen+3]:
            self.checksum += ord(c)

        self.content = rawcontent[3:self.reclen+2]

    def dump(self, out):
        out.write("{0} (0x{1:02x}) record ({2} bytes)\n".format(self.rectypestr, self.rectype, self.reclen))

class OMF_COMENT(OMF_Record):
    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "COMENT"
        self.comment_type = ord(self.content[0])
        self.comment_class = ord(self.content[1])

    def dump(self, out):
        OMF_Record.dump(self, out)
        out.write("  comment type: 0x{0:02x}\n".format(self.comment_type))
        out.write("  comment class: 0x{0:02x}\n".format(self.comment_class))

class OMF_MODEND(OMF_Record):
    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "MODEND"
        self.module_type = ord(self.content[0])
        self.module_type_mattr_main = (self.module_type & 0x80) >> 7
        self.module_type_mattr_start = (self.module_type & 0x40) >> 6
        self.module_type_segment_bit = (self.module_type & 0x20) >> 5
        self.module_type_x = self.module_type & 0x01

        if self.module_type_mattr_start == 1:
            self.end_data = ord(self.content[1])
            self.end_data_f = (self.end_data & 0x80) >> 7
            self.end_data_frame = (self.end_data & 0x70) >> 4
            self.end_data_t = (self.end_data & 0x08) >> 3
            self.end_data_p = (self.end_data & 0x04) >> 2
            self.end_data_targt = self.end_data & 0x03

            offset = 2

            # F = 0 means that frame method and datum are specified
            # within this record (and that no possibly existing
            # frame thread is referred to); as an additional
            # complication, only frame methods 0 to 2 actually need a datum
            if self.end_data_f == 0 and self.end_data_frame >= 0 and self.end_data_frame <= 2:
                self.frame_datum = ord(self.content[offset])
                if (self.frame_datum & 0x80) != 0:
                    self.frame_datum = ((self.frame_datum & 0x7f) << 8) | ord(self.content[offset + 1])
                    offset += 2
                else:
                    offset += 1
            else:
                self.frame_datum = None

            # T = 0 means that target method and datum are specified
            # within this record (and that no possibly existing
            # target thread is referred to)
            if self.end_data_t == 0:
                self.target_datum = ord(self.content[offset])
                if (self.target_datum & 0x80) != 0:
                    self.target_datum = ((self.target_datum & 0x7f) << 8) | ord(self.content[offset + 1])
                    offset += 2
                else:
                    offset += 1

            if self.end_data_p == 0:
                self.target_displacement = ord(self.content[offset + 0]) | (ord(self.content[offset + 1]) << 8)

    def dump(self, out):
        OMF_Record.dump(self, out)

        out.write("  module_type: 0x{0:02x}\n".format(self.module_type))
        out.write("  module_type_mattr_main: 0x{0:02x}\n".format(self.module_type_mattr_main))
        out.write("  module_type_mattr_start: 0x{0:02x}\n".format(self.module_type_mattr_start))
        out.write("  module_type_segment_bit: 0x{0:02x}\n".format(self.module_type_segment_bit))
        out.write("  module_type_x: 0x{0:02x}\n".format(self.module_type_x))

        if self.module_type_mattr_start == 1:
            out.write("  end_data: 0x{0:02x}\n".format(self.end_data))
            out.write("  end_data_f: 0x{0:02x}\n".format(self.end_data_f))
            out.write("  end_data_frame: 0x{0:02x}\n".format(self.end_data_frame))
            out.write("  end_data_t: 0x{0:02x}\n".format(self.end_data_t))
            out.write("  end_data_p: 0x{0:02x}\n".format(self.end_data_p))
            out.write("  end_data_targt: 0x{0:02x}\n".format(self.end_data_targt))
            if self.frame_datum is not None:
                out.write("  frame_datum: 0x{0:04x}\n".format(self.frame_datum))
            out.write("  target_datum: 0x{0:04x}\n".format(self.target_datum))
            if self.end_data_p == 0:
                out.write("  target_displacement: 0x{0:04x}\n".format(self.target_displacement))

# Base type for EXTDEF and LEXTDEF which have exactly the same format
class OMF_EXTDEF_base(OMF_Record):

    class entry(object):
        pass

    def __init__(self, rawcontent, rectypestr):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = rectypestr

        self.entries = []
        offset = 0

        while offset < len(self.content):

            entry = self.entry()

            entry.strlen = ord(self.content[offset + 0])

            if offset + 1 + entry.strlen >= len(self.content):
                raise RuntimeError(self.rectypestr)

            entry.strcontent = self.content[offset + 1:offset  + 1 + entry.strlen]

            entry.index = ord(self.content[offset + 1 + entry.strlen])
            if (entry.index & 0x80) != 0:
                entry.index = ((entry.index & 0x7f) << 8) | ord(self.content[offset + 1 + self.strlen + 1])
                offset += 1 + entry.strlen + 2
            else:
                offset += 1 + entry.strlen + 1

            self.entries.append(entry)

    def dump(self, out):
        OMF_Record.dump(self, out)
        for entry in self.entries:
            out.write("  entry:\n")
            out.write("    string length: {0}\n".format(entry.strlen))
            out.write("    external name string: '{0}'\n".format(entry.strcontent))
            out.write("    index: 0x{0:02x}\n".format(entry.index))

class OMF_EXTDEF(OMF_EXTDEF_base):
    def __init__(self, rawcontent):
        OMF_EXTDEF_base.__init__(self, rawcontent, rectypestr="EXTDEF")

class OMF_LEXTDEF(OMF_EXTDEF_base):
    def __init__(self, rawcontent):
        OMF_EXTDEF_base.__init__(self, rawcontent, rectypestr="LEXTDEF")

# Base type for PUBDEF and LPUBDEF which have exactly the same format
class OMF_PUBDEF_base(OMF_Record):

    class entry(object):
        pass

    def __init__(self, rawcontent, rectypestr):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = rectypestr

        offset = 0

        self.base_group_index = ord(self.content[offset])
        if (self.base_group_index & 0x80) != 0:
            self.base_group_index = ((self.base_group_index & 0x7f) << 8) | ord(self.content[offset + 1])
            offset += 2
        else:
            offset += 1

        self.base_segment_index = ord(self.content[offset])
        if (self.base_segment_index & 0x80) != 0:
            self.base_segment_index = ((self.base_segment_index & 0x7f) << 8) | ord(self.content[offset + 1])
            offset += 2
        else:
            offset += 1

        if self.base_segment_index == 0:
            self.base_frame = ord(self.content[offset + 0]) | (ord(self.content[offset + 1]) << 8)
            offset += 2
        else:
            self.base_frame = 0

        self.entries = []

        while offset < (len(self.content) - 4):
            entry = self.entry()

            entry.strlen = ord(self.content[offset])
            entry.strcontent = self.content[offset+1:offset+1+entry.strlen]
            offset += entry.strlen + 1

            entry.public_offset = ord(self.content[offset + 0]) | (ord(self.content[offset + 1]) << 8)
            offset += 2

            entry.type_index = ord(self.content[offset])
            if (entry.type_index & 0x80) != 0:
                entry.type_index = ((entry.type_index & 0x7f) << 8) | ord(self.content[offset + 1])
                offset += 2
            else:
                offset += 1

            self.entries.append(entry)

        if offset != len(self.content):
            raise RuntimeError("Failed to completetly parse {0} record ({1} bytes remaining)\n".format(self.rectypestr, len(self.content) - offset))

    def dump(self, out):
        OMF_Record.dump(self, out)

        out.write("  base_group_index: 0x{0:04x}\n".format(self.base_group_index))
        out.write("  base_segment_index: 0x{0:04x}\n".format(self.base_segment_index))
        out.write("  base_frame: 0x{0:04x}\n".format(self.base_frame))

        for entry in self.entries:
            out.write("  entry:\n")
            out.write("    string len: {0}\n".format(entry.strlen))
            out.write("    string content: '{0}'\n".format(entry.strcontent))
            out.write("    public offset: 0x{0:04x}\n".format(entry.public_offset))
            out.write("    type index: 0x{0:04x}\n".format(entry.type_index))

class OMF_PUBDEF(OMF_PUBDEF_base):
    def __init__(self, rawcontent):
        OMF_PUBDEF_base.__init__(self, rawcontent, rectypestr="PUBDEF")

class OMF_LPUBDEF(OMF_PUBDEF_base):
    def __init__(self, rawcontent):
        OMF_PUBDEF_base.__init__(self, rawcontent, rectypestr="LPUBDEF")

class OMF_LNAMES(OMF_Record):

    class entry(object):
        pass

    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "LNAMES"

        contentlen = len(self.content)
        offset = 0

        self.entries = []

        while offset < contentlen:

            entry = self.entry()

            entry.strlen = ord(self.content[offset])
            entry.strcontent = self.content[offset+1:offset+1+entry.strlen]
            offset += entry.strlen + 1

            self.entries.append(entry)

    def dump(self, out):
        OMF_Record.dump(self, out)
        for entry in self.entries:
            out.write("  string len: {0}\n".format(entry.strlen))
            out.write("  string content: '{0}'\n".format(entry.strcontent))

class OMF_SEGDEF(OMF_Record):
    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "SEGDEF"

        self.segattr_alignment   = (ord(self.content[0]) & 0xe0) >> 5
        self.segattr_combination = (ord(self.content[0]) & 0x1c) >> 2
        self.segattr_big         = (ord(self.content[0]) & 0x02) >> 1
        self.segattr_p           =  ord(self.content[0]) & 0x01

        if self.segattr_alignment == 0:
            self.segattr_frame_number = ord(self.content[1]) | (ord(self.content[2]) << 8)
            self.segattr_offset = ord(self.content[3])
            idx = 4
        else:
            idx = 1

        self.segment_length = ord(self.content[idx + 0]) | (ord(self.content[idx + 1]) << 8)
        idx += 2

        self.segment_name_index = ord(self.content[idx + 0])
        self.class_name_index   = ord(self.content[idx + 1])
        self.overlay_name_index = ord(self.content[idx + 2])

    def dump(self, out):
        OMF_Record.dump(self, out)
        out.write("  segment alignment:      0x{0:02x}\n".format(self.segattr_alignment))
        out.write("  segment combination:    0x{0:02x}\n".format(self.segattr_combination))
        out.write("  segment big:            0x{0:02x}\n".format(self.segattr_big))
        out.write("  segment p:              0x{0:02x}\n".format(self.segattr_p))
        if self.segattr_alignment == 0:
            out.write("  segment frame number: 0x{0:04x}\n".format(self.segattr_frame_number))
            out.write("  segment offset:         0x{0:02x}\n".format(self.segattr_offset))
        out.write("  segment length:       0x{0:04x}\n".format(self.segment_length))
        out.write("  segment_name_index:     0x{0:02x}\n".format(self.segment_name_index))
        out.write("  class_name_index:       0x{0:02x}\n".format(self.class_name_index))
        out.write("  overlay_name_index:     0x{0:02x}\n".format(self.overlay_name_index))

class OMF_GRPDEF(OMF_Record):

    class entry(object):
        pass

    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "GRPDEF"

        offset = 0

        self.group_name_index = ord(self.content[offset])
        if (self.group_name_index & 0x80) != 0:
            self.group_name_index = ((self.group_name_index & 0x7f) << 8) | ord(self.content[offset + 1])
            offset += 2
        else:
            offset += 1

        self.entries = []

        while (len(self.content) - offset) >= 2:

            entry = self.entry()

            entry.type = ord(self.content[offset])
            offset += 1

            entry.segment_definition = ord(self.content[offset])
            if (entry.segment_definition & 0x80) != 0:
                entry.segment_definition = ((entry.segment_definition & 0x7f) << 8) | ord(self.content[offset + 1])
                offset += 2
            else:
                offset += 1

            self.entries.append(entry)

        if offset != len(self.content):
            raise RuntimeError("Failed to completetly parse GRPDEF record ({0} bytes remaining)".format(len(self.content) - offset))

    def dump(self, out):
        OMF_Record.dump(self, out)

        out.write("  group_name_index: 0x{0:02x}\n".format(self.group_name_index))

        for entry in self.entries:
            out.write("  entry:\n")
            out.write("    type: 0x{0:02x}\n".format(entry.type))
            out.write("    segment_definition: 0x{0:02x}\n".format(entry.segment_definition))

class OMF_FIXUPP(OMF_Record):

    class THREAD_subrecord(object):
        pass

    class FIXUP_subrecord(object):
        pass

    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "FIXUPP"

        offset = 0

        self.subrecords = []

        while offset < len(self.content):
            b = ord(self.content[offset])

            if (b & 0x80) == 0:
                subrecord = self.THREAD_subrecord()

                subrecord.thread_d = (b & 0x40) >> 6
                subrecord.thread_method = (b & 0x1c) >> 2
                subrecord.thread_thred = b & 0x03
                offset += 1

                if subrecord.thread_d == 0 and (subrecord.thread_method & 0x4) != 0:
                    raise RuntimeError("THREAD subrecord defines a TARGET thread and has bit 0x4 set!?")

                if subrecord.thread_method >= 0 and subrecord.thread_method <= 2:
                    subrecord.index = ord(self.content[offset])
                    if (subrecord.index & 0x80) != 0:
                        subrecord.index = ((subrecord.index & 0x7f) << 8) | ord(self.content[offset + 1])
                        offset += 2
                    else:
                        offset += 1
            else:
                subrecord = self.FIXUP_subrecord()

                subrecord.fixup_m = (b & 0x40) >> 6
                subrecord.fixup_location = (b & 0x3c) >> 2
                subrecord.fixup_data_record_offset = ((b & 0x03) << 8) | ord(self.content[offset + 1])
                offset += 2

                subrecord.fix_data = ord(self.content[offset])
                subrecord.fix_data_f = (subrecord.fix_data & 0x80) >> 7
                subrecord.fix_data_frame = (subrecord.fix_data & 0x70) >> 4
                subrecord.fix_data_t = (subrecord.fix_data & 0x08) >> 3
                subrecord.fix_data_p = (subrecord.fix_data & 0x04) >> 2
                subrecord.fix_data_targt = subrecord.fix_data & 0x03
                offset += 1

                # F = 0 means that frame method and datum are specified
                # within this FIXUP subrecord (and that no possibly existing
                # frame thread is referred to); as an additional
                # complication, only frame methods 0 to 2 actually need a datum
                if subrecord.fix_data_f == 0 and subrecord.fix_data_frame >= 0 and subrecord.fix_data_frame <= 2:
                    subrecord.frame_datum = ord(self.content[offset])
                    if (subrecord.frame_datum & 0x80) != 0:
                        subrecord.frame_datum = ((subrecord.frame_datum & 0x7f) << 8) | ord(self.content[offset + 1])
                        offset += 2
                    else:
                        offset += 1
                else:
                    subrecord.frame_datum = None

                # T = 0 means that target method and datum are specified
                # within this FIXUP subrecord (and that no possibly existing
                # target thread is referred to)
                if subrecord.fix_data_t == 0:
                    subrecord.target_datum = ord(self.content[offset])
                    if (subrecord.target_datum & 0x80) != 0:
                        subrecord.target_datum = ((subrecord.target_datum & 0x7f) << 8) | ord(self.content[offset + 1])
                        offset += 2
                    else:
                        offset += 1

                if subrecord.fix_data_p == 0:
                    subrecord.target_displacement = ord(self.content[offset + 0]) | (ord(self.content[offset + 1]) << 8)
                    offset += 2

            self.subrecords.append(subrecord)

    def dump(self, out):
        OMF_Record.dump(self, out)
        for subrecord in self.subrecords:
            if isinstance(subrecord, self.THREAD_subrecord):
                out.write("  THREAD subrecord:")
                out.write("    thread_d: {0:02x}\n".format(subrecord.thread_d))
                out.write("    thread_method: {0:02x}\n".format(subrecord.thread_method))
                out.write("    thread_thred: {0:02x}\n".format(subrecord.thread_thred))
                if subrecord.thread_method >= 0 and subrecord.thread_method <= 2:
                    out.write("    index: {0:04x}\n".format(subrecord.index))
            elif isinstance(subrecord, self.FIXUP_subrecord):
                out.write("  FIXUP subrecord:\n")
                out.write("    fixup_m: 0x{0:02x}\n".format(subrecord.fixup_m))
                out.write("    fixup_location: 0x{0:02x}\n".format(subrecord.fixup_location))
                out.write("    fixup_data_record_offset: 0x{0:04x}\n".format(subrecord.fixup_data_record_offset))
                out.write("    fix_data: 0x{0:02x}\n".format(subrecord.fix_data))
                out.write("    fix_data_f: 0x{0:02x}\n".format(subrecord.fix_data_f))
                out.write("    fix_data_frame: 0x{0:02x}\n".format(subrecord.fix_data_frame))
                out.write("    fix_data_t: 0x{0:02x}\n".format(subrecord.fix_data_t))
                out.write("    fix_data_p: 0x{0:02x}\n".format(subrecord.fix_data_p))
                out.write("    fix_data_targt: 0x{0:02x}\n".format(subrecord.fix_data_targt))
                if subrecord.fix_data_f == 0 and subrecord.fix_data_frame >= 0 and subrecord.fix_data_frame <= 2:
                    out.write("    frame_datum: 0x{0:04x}\n".format(subrecord.frame_datum))
                out.write("    target_datum: 0x{0:04x}\n".format(subrecord.target_datum))
                if subrecord.fix_data_p == 0:
                    out.write("    target_displacement: 0x{0:04x}\n".format(subrecord.target_displacement))

# Base type for LEDATA and LIDATA which we treat equally while parsing
# (LIDATA can't be expanded yet because FIXUPs need to be applied before)
class OMF_LEDATA_base(OMF_Record):
    def __init__(self, rawcontent, rectypestr):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = rectypestr
        self.segment_index = ord(self.content[0])
        self.data_offset = ord(self.content[1]) | (ord(self.content[2]) << 8)
        self.data = self.content[3:]
        self.hexdump = ""
        for b in self.data:
            self.hexdump += "{0:02x}".format(ord(b))

    def dump(self, out):
        OMF_Record.dump(self, out)
        out.write("  segment index: {0}\n".format(self.segment_index))
        out.write("  data offset: {0}\n".format(self.data_offset))
        out.write("  data: {0}\n".format(self.hexdump))

class OMF_LEDATA(OMF_LEDATA_base):
    def __init__(self, rawcontent):
        OMF_LEDATA_base.__init__(self, rawcontent, rectypestr="LEDATA")

class OMF_LIDATA(OMF_LEDATA_base):
    def __init__(self, rawcontent):
        OMF_LEDATA_base.__init__(self, rawcontent, rectypestr="LIDATA")

    # Expands the iterated data (must be done *after* FIXUPs are applied)
    def get_expanded_data(self):
        offset = 3
        data = ""
        while offset < len(self.content):
            if (len(self.content) - offset) < 6:
                raise RuntimeError("Got LIDATA record with data block smaller than 6 bytes which is invalid")
            repeat_count = ord(self.content[offset + 0]) | (ord(self.content[offset + 1]) << 8)
            block_count = ord(self.content[offset + 2]) | (ord(self.content[offset + 3]) << 8)
            if block_count != 0:
                # We don't support recursive definitions as the OpenWatcom compiler doesn't seem to emit them anyway
                raise RuntimeError("Got LIDATA record data block with block count != 0 which is unsupported")
            count = ord(self.content[offset + 4])
            content = self.content[offset + 5:offset + 5 + count]
            data += content * block_count
            offset += 4 + count

class OMF_THEADR(OMF_Record):
    def __init__(self, rawcontent):
        OMF_Record.__init__(self, rawcontent)
        self.rectypestr = "THEADR"
        self.strlen = ord(self.content[0])
        self.strcontent = self.content[1:]
        if len(self.content) != len(self.strcontent) + 1:
            raise RuntimeError("Inconsistent string length ({0}/{1})\n".format(len(self.content), len(strcontent) + 2))

    def dump(self, out):
        OMF_Record.dump(self, out)
        out.write("  string len: {0}\n".format(self.strlen))
        out.write("  string content: '{0}'\n".format(self.strcontent))

def get_next_record(data, offset):
    rectype = ord(data[offset + 0])
    reclen = ord(data[offset + 1]) | (ord(data[offset + 2]) << 8)

    if rectype == 0x80:
        obj = OMF_THEADR(data[offset:offset+3+reclen])
    elif rectype == 0x88:
        obj = OMF_COMENT(data[offset:offset+3+reclen])
    elif rectype == 0x8a:
        obj = OMF_MODEND(data[offset:offset+3+reclen])
    elif rectype == 0x8c:
        obj = OMF_EXTDEF(data[offset:offset+3+reclen])
    elif rectype == 0x90:
        obj = OMF_PUBDEF(data[offset:offset+3+reclen])
    elif rectype == 0x96:
        obj = OMF_LNAMES(data[offset:offset+3+reclen])
    elif rectype == 0x98:
        obj = OMF_SEGDEF(data[offset:offset+3+reclen])
    elif rectype == 0x9a:
        obj = OMF_GRPDEF(data[offset:offset+3+reclen])
    elif rectype == 0x9c:
        obj = OMF_FIXUPP(data[offset:offset+3+reclen])
    elif rectype == 0xa0:
        obj = OMF_LEDATA(data[offset:offset+3+reclen])
    elif rectype == 0xa2:
        obj = OMF_LIDATA(data[offset:offset+3+reclen])
    elif rectype == 0xb4:
        obj = OMF_LEXTDEF(data[offset:offset+3+reclen])
    elif rectype == 0xb6:
        obj = OMF_LPUBDEF(data[offset:offset+3+reclen])
    else:
        raise RuntimeError("Unknown record type 0x{0:02x}".format(rectype))

    return obj

